#!/usr/bin/perl
# vim: set noet ts=4 sw=4

use strict;
use Getopt::Std;
use Net::DNS::Nameserver;
use Net::DNS::Resolver;
use LWP::UserAgent;
use HTTP::Request::Common;
use Sys::Syslog qw(:standard :macros);
use Data::Dumper;

require '/etc/webnamed.conf';

our ($proxy, %rmap, @resolvers, $listen_addr, $listen_port,
	$domain_name, $agent, $forward_to_default_dns,
	$debug, $pkt_dump);

my $command_name = 'webnamed';
my $VERSION = '0.3.1';
$0 = $command_name;

my $resolver_idx = 0;
my %amap = ();

sub plog {
	my ($prio, $s) = @_;

	if ($debug) {
		print "$s\n";
	} else {
		syslog($prio, $s);
	}
}

sub create_req {
	my ($qname, $resolver) = @_;
	my $req = GET("$resolver?q=$qname");
	$req;
}

sub pickup_addr {
	my ($content) = @_;
	if ($$content =~ /^result:([0-9\.]+) ([0-9]+)/m) {
		($1, $2);
	} else {
		(undef, 0);
	}
}

sub resolve_by_outer_service {

	my ($qname, $qclass, $qtype, $peerhost, $query, $conn) = @_;
	my ($rcode, @ans, $ttl, $addr, $cache);

	$cache = $amap{$qname};
	if ($cache ne undef) {
		my ($tm, $rtm);
		($addr, $ttl, $tm) = @$cache;
		$rtm = time() - $tm;
		if ($rtm > $ttl) {
			$addr = $amap{$qname} = undef;
		} else {
			$ttl = $ttl - $rtm;
			$rcode = "NOERROR";
			if ($debug) {
				print "#Hit Cache: $qname -> $addr\n";
			}
		}
	}

	if ($addr eq undef) {
		my $res;

		for (my $i = 0; $i < @resolvers; ++$i) {
			my $resolver_type = $resolvers[$resolver_idx];
			my $referer = $rmap{$resolver_type}->{referer};
			my $resolver = $rmap{$resolver_type}->{resolver};
			my ($ua, $req);

			if ($resolver_idx == @resolvers - 1) {
				$resolver_idx = 0;
			} else {
				$resolver_idx++;
			}
			if ($debug) {
				print "#Request to $resolver\n";
			}
			$ua = LWP::UserAgent->new;
			$ua->agent($agent);
			$req = create_req($qname, $resolver);
			$req->referer($referer);
			if (length($proxy) > 0) {
				$ua->proxy('http', $proxy);
			}
			$res = $ua->request($req);
			last if ($res->is_success);
			if ($debug) {
				print "#Error Response\n";
				print $res->error_as_HTML;
				print "\n".("-" x 16)."\n";
			}
			plog(LOG_ERR, "resolver [$resolver_type] ".
				"returns error ".($res->code).": ".($res->message));
		}

		$ttl = 0;
		if ($res->is_success) {
			($addr, $ttl) = pickup_addr(\$res->content);
			if ($addr eq undef) {
				$rcode = "DXDOMAIN";
				if ($debug) {
					print "#Unresolved: $qname\n";
					print $res->content;
					print "\n".("-" x 16)."\n";
				}
			} else {
				$rcode = "NOERROR";
				$amap{$qname} = [ $addr, $ttl, time() ];
				if ($debug) {
					print "#Resolved: $qname -> $addr\n";
				}
			}
		} else {
			$rcode = "DXDOMAIN";
		}
	}

	push @ans, Net::DNS::RR->new("$qname $ttl $qclass $qtype $addr");
	($rcode, \@ans);
}

sub resolve_by_dns {

	my ($qname, $qclass, $qtype, $peerhost, $query, $conn) = @_;
	my ($rcode, @ans, $addr, $ttl);

	if ($debug) {
		print "\n#Resolving by default DNS: $qname $qtype\n";
	}

	my $res = Net::DNS::Resolver->new;
	my $pkt = $res->search($qname, $qtype);

	if ($pkt eq undef) {
		$rcode = "NXDOMAIN";
		$ttl = 0;
		if ($debug) {
			print "#Unresolved (pkt is null)\n";
		}
	} else {
		my @dnsans = $pkt->answer;
		if ($pkt_dump) {
			print Data::Dumper->Dump([\@dnsans], [qw(*dnsans)]);
		}
		my $idx = (@dnsans > 1) ? 1 : 0;
		$addr = $dnsans[$idx]->{address};
		$ttl = $dnsans[$idx]->{ttl};
		if (($addr eq undef) || length($addr) eq 0) {
			$rcode = "NXDOMAIN";
			if ($debug) {
				print "#Unresolved (no addr)\n";
			}
		} else {
			$rcode = "NOERROR";
			$amap{$qname} = [ $addr, $ttl, time() ];
			if ($debug) {
				print "#Resolved: $addr\n";
			}

		}
	}
	push @ans, Net::DNS::RR->new("$qname $ttl $qclass $qtype $addr");
	($rcode, \@ans);
}

sub query_reply_handler {

	my ($qname, $qclass, $qtype, $peerhost, $query, $conn) = @_;
	my ($rcode, $ans, @auth, @add);

	if ($debug) {
		print "\n#Request from $peerhost to ".($conn->{"sockhost"}).
			" qname=$qname qtype=$qtype\n";
	}

	if ($qtype eq 'A' && $qname !~ /\.${domain_name}$/) {
		($rcode, $ans) = resolve_by_outer_service($qname, $qclass, $qtype,
			$peerhost, $query, $conn);
	}
	if (($rcode ne "NOERROR") && $forward_to_default_dns) {
		($rcode, $ans) = resolve_by_dns($qname, $qclass, $qtype,
			$peerhost, $query, $conn);
	}

	return ($rcode, $ans, \@auth, \@add, { aa => 1 });
}

sub dump_amap {
	my ($sig) = @_;

	$SIG{$sig} = 'IGNORE';

	plog(LOG_DEBUG, "#dump cache : BEGIN");

	my $ctm = time();
	foreach my $qname (keys %amap) {
		my $cache = $amap{$qname};
		my $s = "# $qname ";
		if ($cache eq undef) {
			$s .= "error";
		} else {
			my ($addr, $ttl, $tm) = @$cache;
			my $ttl = $ttl - ($ctm - $tm);
			if ($ttl <= 0) {
				$ttl = "expired";
			}
			$s .= "$addr $ttl";
		}
		plog(LOG_DEBUG, "$s");
	}

	plog(LOG_DEBUG, "#dump cache : END");

	$SIG{$sig} = (caller 0)[3];
}

sub shutdown {
	my ($sig) = @_;

	$SIG{$sig} = 'IGNORE';
	plog(LOG_INFO, "shutdown");
	closelog();
	print "#Shutdown\n" if ($debug);
	exit(0);
}

sub psuedo_daemonize {
	my $r = fork();
	if ($r eq undef) {
		die "Failed to fork process: $!\n";
	}
	if ($r > 0) {
		exit(0);
	}

	my $pidfile = "/var/run/$command_name.pid"; 
	open(OUT, ">$pidfile") || die "Failed to write $pidfile: $!";
	print OUT $$;
	close OUT;
}

sub run {
	my %opts;

	getopts("hd", \%opts);
	foreach my $opt (keys %opts) {
		if ($opt eq 'h') {
			usage();
		} elsif ($opt eq 'd') {
			$debug = 1;
		}
	}

	my $ns = Net::DNS::Nameserver->new(
		LocalAddr => $listen_addr,
		LocalPort => $listen_port,
		ReplyHandler => \&query_reply_handler,
		Verbose => $debug
	);
	eval {
		$ns->loop_once(0);
	};
	if ($@) {
		print "Cannot start to listen socket at ".
				"$listen_addr:$listen_port.\n";
		exit 1;
	}

	openlog($command_name, 'pid', 'local0')
		|| die "Failed to open syslog: $!\n";

	$SIG{USR1} = 'dump_amap';
	$SIG{INT} = 'shutdown';
	$SIG{TERM} = 'shutdown';

	if (!$debug) {
		psuedo_daemonize();
	}

	plog(LOG_INFO, "start");
	$ns->main_loop;
}

sub usage {
	print <<_EOL_;
$command_name $VERSION

Usage: $command_name [-h] [-d]

  -h : Show this message.
  -d : Run as a foreground process and output debug log.

Copyright (C) 2011 t.baba All rights reserved.
Licensed under the GPLv2.
_EOL_
	exit 1;
}

run();
